package my.shadowsocks.client;

import io.netty.bootstrap.Bootstrap;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.*;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.socks.SocksAddressType;
import io.netty.handler.codec.socks.SocksCmdRequest;
import io.netty.handler.codec.socks.SocksCmdResponse;
import io.netty.handler.codec.socks.SocksCmdStatus;
import io.netty.util.concurrent.FutureListener;
import io.netty.util.concurrent.Promise;
import my.shadowsocks.common.config.Config;
import my.shadowsocks.common.encryption.Crypt;
import my.shadowsocks.common.encryption.CryptFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;

/**
 * @author fengbo
 * @date 2018/10/8
 */
public class SocksServerConnectHandler extends SimpleChannelInboundHandler<SocksCmdRequest> {

    private static final Logger log = LoggerFactory.getLogger(SocksServerConnectHandler.class);

    private static final Set<String> PAC_SET = new HashSet<>(20);

    static {
        PAC_SET.add("youtube.com");
        PAC_SET.add("ytimg.com");
        PAC_SET.add("ggpht.com");
        PAC_SET.add("google.com");
        PAC_SET.add("google.co.jp");
        PAC_SET.add("googlevideo.com");
        PAC_SET.add("googleapis.com");
        PAC_SET.add("googleusercontent.com");
        PAC_SET.add("wikipedia.org");
        PAC_SET.add("shadowsocks.org");
        PAC_SET.add("facebook.net");
        PAC_SET.add("twitter.com");
        PAC_SET.add("quora.com");
    }

    private final Bootstrap b = new Bootstrap();
    private boolean isProxy = true;
    private ChannelHandlerContext ctx;
    private Channel outboundChannel;
    private final Crypt crypt;

    SocksServerConnectHandler() {
        Config config = Config.getInstance();
        crypt = CryptFactory.get(config.getStringValue("method"), config.getStringValue("password"));
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, SocksCmdRequest request) {

        Promise<Channel> promise = ctx.executor().newPromise();
        promise.addListener((FutureListener<Channel>) future -> {
            outboundChannel = future.getNow();
            if (future.isSuccess()) {
                AbstractRelayHandler inRelay = new InAbstractRelayHandler(SocksServerConnectHandler.this);
                AbstractRelayHandler outRelay = new OutAbstractRelayHandler(SocksServerConnectHandler.this);
                if (isProxy) {
                    sendConnectRemoteMessage(request);
                }

                ctx.channel().writeAndFlush(new SocksCmdResponse(SocksCmdStatus.SUCCESS, SocksAddressType.IPv4))
                        .addListener((ChannelFutureListener) channelFuture -> {
                    ctx.pipeline().remove(SocksServerConnectHandler.this);
                    outboundChannel.pipeline().addLast(inRelay);
                    ctx.pipeline().addLast(outRelay);
                });
            } else {
                ctx.channel().writeAndFlush(new SocksCmdResponse(SocksCmdStatus.FAILURE, SocksAddressType.IPv4));
                SocksServerUtils.closeOnFlush(ctx.channel());
            }
        });

        b.group(ctx.channel().eventLoop()).channel(NioSocketChannel.class)
                .option(ChannelOption.CONNECT_TIMEOUT_MILLIS, 10000)
                .option(ChannelOption.SO_KEEPALIVE, true)
                .handler(new DirectClientHandler(promise));

        String host = request.host();
        int port = request.port();
//        setProxy(host);

        log.info("host = {}, port = {}, isProxy = {}", host, port, isProxy);

        b.connect(getHost(host), getPort(port)).addListener((ChannelFutureListener) future -> {
            if (!future.isSuccess()) {
                ctx.channel().writeAndFlush(new SocksCmdResponse(SocksCmdStatus.FAILURE, SocksAddressType.IPv4));
                SocksServerUtils.closeOnFlush(ctx.channel());
            }
        });
    }

    @Override
    public void handlerAdded(ChannelHandlerContext ctx) {
        this.ctx = ctx;
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
        log.error("SocksServerConnectHandler throw error!", cause);
        SocksServerUtils.closeOnFlush(ctx.channel());
    }

    private void setProxy(String host) {
        if (!isProxy(host)) {
            isProxy = false;
        }
    }

    /**
     * 判断是否需要被代理
     * @param host 主机名
     * @return 是否需要被代理
     */
    private boolean isProxy(String host) {
        String suffix;
        int pos = host.lastIndexOf('.');
        pos = host.lastIndexOf('.', pos - 1);
        while (true) {
            if (pos <= 0) {
                return PAC_SET.contains(host);
            }
            suffix = host.substring(pos + 1);
            if (PAC_SET.contains(suffix)) {
                return true;
            }
            pos = host.lastIndexOf('.', pos - 1);
        }
    }

    /**
     * 获取真正访问的主机地址
     * @param defaultVal 默认值
     * @return 结果
     */
    private String getHost(String defaultVal) {
        if (isProxy) {
            return Config.getInstance().getStringValue("remote.address");
        }
        return defaultVal;
    }

    /**
     * 获取真正访问的端口号
     * @param defaultVal 默认值
     * @return 结果
     */
    private int getPort(int defaultVal) {
        if (isProxy) {
            return Config.getInstance().getIntValue("remote.port");
        }
        return defaultVal;
    }

    private void sendConnectRemoteMessage(SocksCmdRequest request) {
        ByteBuf buff = Unpooled.buffer();
        request.encodeAsByteBuf(buff);
        if (buff.hasArray()) {
            int len = buff.readableBytes();
            byte[] arr = new byte[len];
            buff.getBytes(0, arr);
            byte[] data = remoteByte(arr);
            sendRemote(data, data.length);
        }
    }

    /**
     * 发送给远程的数据跳过3个字节
     * @param data 被代理客户端传递过来的数据
     * @return 将要发送给远程的数据
     */
    private byte[] remoteByte(byte[] data) {
        int dataLength = data.length;
        dataLength -= 3;
        byte[] temp = new byte[dataLength];
        System.arraycopy(data, 3, temp, 0, dataLength);
        return temp;
    }

    /**
     * 向远程代理服务器发送数据
     * @param data 数据
     * @param length 数据长度
     */
    void sendRemote(byte[] data, int length) {
        ByteArrayOutputStream remoteOutStream = new ByteArrayOutputStream();
        try {
            if (isProxy) {
                crypt.encrypt(data, length, remoteOutStream);
                data = remoteOutStream.toByteArray();
            }
            outboundChannel.writeAndFlush(Unpooled.wrappedBuffer(data));
        } catch (Exception e) {
            log.error("send remote error", e);
        } finally {
            try {
                remoteOutStream.close();
            } catch (IOException e) {
                log.error("close remote stream error", e);
            }
        }
    }

    /**
     * 给被代理的客户端发送访问结果
     * @param data 发送的数据
     * @param length 数据长度
     */
    void sendLocal(byte[] data, int length) {
        ByteArrayOutputStream localOutStream = new ByteArrayOutputStream();
        try {
            if (isProxy) {
                crypt.decrypt(data, length, localOutStream);
                data = localOutStream.toByteArray();
            }
            ctx.writeAndFlush(Unpooled.wrappedBuffer(data));
        } catch (Exception e) {
            log.error("send local error", e);
        } finally {
            try {
                localOutStream.close();
            } catch (IOException e) {
                log.error("close local stream error", e);
            }
        }
    }
}
